import math
import pandas as pd


def create_new_data(new_head, new_lines):  # 构建dataframe
    t = {}
    index = 0
    for i in new_head:
        line = []
        for j in new_lines:
            line.append(j[index])
        t[i] = line
        index += 1
    return pd.DataFrame(t)


def getinfo_d(t, label):  # 计算D
    count_yes = len(t[t[label] == 'yes'])
    count_no = len(t[t[label] == 'no'])

    d = get_d(count_yes, count_no)

    return round(d, 3)


# 计算最小d的索引，对应需要剪掉的列索引
def get_min_d_index(head, t, label):  # 获取一个属性和buys_computer的D,求得最小值的索引
    d = []

    for a in head:  # 对每个节点的所有属性进行遍历计算d
        if a == label:
            break

        index = list(set(t[a].tolist()))  # 得到每个列不同的属性

        temp_d = 0
        for b in index:
            yes = len(t[(t[a] == b) & (t[label] == 'yes')])
            no = len(t[(t[a] == b) & (t[label] == 'no')])

            temp_d += (yes + no) / len(t) * get_d(yes, no)

        d.append(temp_d)  # 把小d累加到大d上面

    d_index = d.index(min(d))  # 找到最小的d的索引

    return d_index


# 计算d
def get_d(yes_counts, no_counts):
    if yes_counts == 0 or no_counts == 0:
        return 0
    else:
        lens = yes_counts + no_counts
        return - yes_counts / lens * math.log2(yes_counts / lens) - no_counts / lens * math.log2(no_counts / lens)
